/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

/**
 * @file channel_mixing.cpp
 */

#include "channel_mixing_310b.h"
#include "channel_mixing_310p.h"

extern "C" __global__ __aicore__ void channel_mixing(GM_ADDR x, GM_ADDR h0, GM_ADDR xk, GM_ADDR kw, GM_ADDR vw, GM_ADDR out, GM_ADDR ht, GM_ADDR hs, GM_ADDR workspace, GM_ADDR tiling)
{
    GET_TILING_DATA(tiling_data, tiling);
    TPipe pipe;

    if (TILING_KEY_IS(1)) 
    {
      ChannelMixingKernel310p op;
      // for(int i=0; i<tiling_data.B; i++)
      // {
        // if (i != 0) {
        //     pipe.Init();
        // }
        REGIST_MATMUL_OBJ(&pipe, GetSysWorkSpacePtr(), op.matmulObj1, &tiling_data.cubeTilingData1, op.matmulObj2, &tiling_data.cubeTilingData2);
        
        op.Init(x, h0, xk, kw, vw, out, ht, hs, workspace, tiling_data.cubeTilingData1, tiling_data.cubeTilingData2, 
                0, tiling_data.T, &pipe);
      for(int i=0; i<tiling_data.B; i++)
      {
        op.Process(&pipe);
      }
        // pipe.Reset();
        // pipe.Destroy();
      // }
    }
    else if (TILING_KEY_IS(2)) 
    {
      ChannelMixingKernel310b op;

      for(int i=0; i<tiling_data.B; i++)
      {
        //  if (i != 0) {
        //     pipe.Init();
        // }

        REGIST_MATMUL_OBJ(&pipe, GetSysWorkSpacePtr(), op.matmulObj);
        op.Init(x, h0, xk, kw, vw, out, ht, hs, tiling_data.cubeTilingData1, tiling_data.cubeTilingData2, 
              i, tiling_data.T, &pipe);
        op.Process(&pipe);
        pipe.Reset();
        // pipe.Destroy();
      }
    }
}
#ifndef ASCENDC_CPU_DEBUG
// call of kernel function
void channel_mixing_do(uint32_t blockDim, void* l2ctrl, void* stream, uint8_t* x, uint8_t* h0, uint8_t* xk, uint8_t* kw, uint8_t* vw, uint8_t* out, uint8_t* ht, uint8_t* hs,
    uint8_t* workspace, uint8_t* tiling)
{
    channel_mixing<<<blockDim, l2ctrl, stream>>>(x, h0, xk, kw, vw, out, ht, workspace, tiling);
}
#endif
